package ssau.knyazev.fft.modules;

import java.io.IOException;
import ssau.knyazev.common.CudaCRuntimeException;
import ssau.knyazev.common.UniException;
import ssau.knyazev.fft.common.FFTRuntimeException;
import ssau.knyazev.fft.common.FFTConst;
import ssau.knyazev.modules.CudaCompiler;
import ssau.knyazev.modules.Generator;
import ssau.knyazev.modules.Printer;
import jcuda.Pointer;
import jcuda.Sizeof;
import jcuda.driver.CUdeviceptr;
import jcuda.driver.CUfunction;
import jcuda.driver.JCudaDriver;

public class GPUFFTModule {

	public static float[][] fft(float[][] src) throws IOException, UniException {
		return fft(src, FFTConst.SIMPLE_MODE);
	}
	
	public static float[][] fft(float[][] src, int mode) throws IOException, UniException {
		String moduleName = null;
		CUfunction function = null;
		if (mode == FFTConst.SIMPLE_MODE){
			moduleName = CudaCompiler.preparePtxFile(FFTConst.CUDA_SIMPLE_MODULE_SRC, true);
			function = CudaCompiler.getFunction(FFTConst.CUDA_FFT_FUNCTION_NAME, moduleName);
		} else if (mode == FFTConst.CCOMPLEX_MIRROR_MODE){
			moduleName = CudaCompiler.preparePtxFile(FFTConst.CUDA_SIMPLE_MODULE_SRC, true);
			function = CudaCompiler.getFunction(FFTConst.CUDA_FFT_FUNCTION_NAME, moduleName);
		} else {
			throw new FFTRuntimeException("This mode is not supported " + mode);
		}
		return fft(src, mode, function);
	}

	public static float[][] fft(float[][] src, int mode, CUfunction function) throws IOException, UniException {
		if (mode == FFTConst.SIMPLE_MODE){
			return fftSimple(src, function);
		} else if (mode == FFTConst.CCOMPLEX_MIRROR_MODE){
			return fftGCMirror(src, function);
		} else {
			throw new FFTRuntimeException("This mode is not supported " + mode);
		}
	}

	private static float[][] fftSimple(float[][] src, CUfunction function) throws IOException {
		_inverse(src);
		float[] _src = toVector(src);
		int N = src.length;
		int size = 0;
		for(int x = src.length; x > 1; x = x >> 1){
			size++;
		}
		
		CUdeviceptr srcPtr = new CUdeviceptr();
		JCudaDriver.cuMemAlloc(srcPtr, _src.length * Sizeof.FLOAT);
		JCudaDriver.cuMemcpyHtoD(srcPtr, Pointer.to(_src), _src.length * Sizeof.FLOAT);
		
		Pointer kernelParameters = Pointer.to(
				Pointer.to(new int[] { size }), 
				Pointer.to(new int[] { N }),
				Pointer.to(srcPtr) 
				);
		JCudaDriver.cuLaunchKernel(function, 
				1, 1, 1, // Grid dimension
				16, 16, 1, // Block dimension
				4096*Sizeof.FLOAT, null, // Shared memory size and stream
				kernelParameters, null // Kernel- and extra parameters
		);
		JCudaDriver.cuCtxSynchronize();

		JCudaDriver.cuMemcpyDtoH(Pointer.to(_src), srcPtr, _src.length * Sizeof.FLOAT);
		JCudaDriver.cuMemFree(srcPtr);
		
		return toMatrix(_src, src.length);
	}

	private static float[][] fftGCMirror(float[][] src, CUfunction function) throws IOException {
		__inverse(src);
		int x = src.length;
		int N = src.length;
		int size = 0;
		while (x > 1) {
			size++;
			x = x >> 1;
		}
		float[] _src = toVector(src);
		
		CUdeviceptr srcPtr = new CUdeviceptr();
		JCudaDriver.cuMemAlloc(srcPtr, _src.length * Sizeof.FLOAT);
		JCudaDriver.cuMemcpyHtoD(srcPtr, Pointer.to(_src), _src.length * Sizeof.FLOAT);
		
		Pointer kernelParameters = Pointer.to(
				Pointer.to(new int[] { size }), 
				Pointer.to(new int[] {N}),
				Pointer.to(srcPtr)
				);
		JCudaDriver.cuLaunchKernel(function, 
				1, 1, 1, // Grid dimension
				20, 20, 1, // Block dimension
				4096*Sizeof.FLOAT, null, // Shared memory size and stream
				kernelParameters, null // Kernel- and extra parameters
		);
		JCudaDriver.cuCtxSynchronize();
		JCudaDriver.cuMemcpyDtoH(Pointer.to(_src), srcPtr, _src.length * Sizeof.FLOAT);
		
		JCudaDriver.cuMemFree(srcPtr);
		
		return toMatrix(_src, src.length);
	}
	
	public static float[][] __fft(float[][] a){
		__inverse(a);
		int x = a.length;
		int N = a.length;
		int size = 0;
		while (x > 1) {
			size++;
			x = x >> 1;
		}
		float[] _a = toVector(a);
		_a = __fft__(size, N, _a);
		a = toMatrix(_a, a.length);
		return a;
	}
	
	private static float[] __fft__(int size, int N, float[] a){
		float[] wlist = new float[N*2];
		
		float[] w1 = new float[4];
		float[] w2 = new float[4];
		float[] tmp = new float[4];
		float[] tmp2 = new float[4];
		
		int step = 0;
		int step2 = 0;
		int flag = 0;
		int wiInd = 0;
		int wjInd = 0;
		int dsInd = 0;
		int rInd = 0;
		int rsInd = 0;
		
		float a00a;
		float a00b;
		float a00c;
		float a00d;
		float a01a;
		float a01b;
		float a01c;
		float a01d;
		float a10a;
		float a10b;
		float a10c;
		float a10d;
		float a11a;
		float a11b;
		float a11c;
		float a11d;
		
		int a00ai;
		int a00bi;
		int a00ci;
		int a00di;
		int a01ai;
		int a01bi;
		int a01ci;
		int a01di;
		int a10ai;
		int a10bi;
		int a10ci;
		int a10di;
		int a11ai;
		int a11bi;
		int a11ci;
		int a11di;
		
		
		for (int i = 0; i < N; i++) {
			double power = -2*Math.PI*i/N;
			wlist[2*i] = (float)Math.cos(power);
			wlist[2*i+1] = (float)Math.sin(power);
		}
		
		for (int q = 0; q < size; q++) {
			step = (int) Math.pow(2, q);
			step2 = step * 2;
			for (int down = 0; down < N - step; down++) {
				for (int right = 0; right < N - step; right++) {
					//flag = (down/step) % 2;
					//flag = flag + (right/step) % 2;
					flag = ((down % step2) <= (step2/4)) ? 0 : 1;
					flag = ((right % step2) <= (step2/4)) ? flag : 1;
					if (flag == 0){
						wiInd = (down % step2)*(N/step2) * 2;
						wjInd = (right % step2)*(N/step2) * 2;
						
						w1[0] = wlist[wiInd];
						w1[1] = wlist[wiInd+1];
						w1[2] = 0;
						w1[3] = 0;
						w2[0] = wlist[wjInd];
						w2[1] = 0;
						w2[2] = wlist[wjInd+1];
						w2[3] = 0;
						
						dsInd = down + step;
						rInd = 4*right;
						rsInd = 4*(right+step);
						
						a00ai = down*N*4 + rInd;
						a00bi = a00ai + 1;
						a00ci = a00ai + 2;
						a00di = a00ai + 3;
						a01ai = down*N*4 + rsInd;
						a01bi = a01ai + 1;
						a01ci = a01ai + 2;
						a01di = a01ai + 3;
						a10ai = dsInd*N*4 + rInd;
						a10bi = a10ai + 1;
						a10ci = a10ai + 2;
						a10di = a10ai + 3;
						a11ai = dsInd*N*4 + rsInd;
						a11bi = a11ai + 1;
						a11ci = a11ai + 2;
						a11di = a11ai + 3;
						
						//XXX = XXX.add(a00);
						a00a = a[a00ai];
						a00b = a[a00bi];
						a00c = a[a00ci];
						a00d = a[a00di];
						a10a = a[a00ai];
						a10b = a[a00bi];
						a10c = a[a00ci];
						a10d = a[a00di];
						a01a = a[a00ai];
						a01b = a[a00bi];
						a01c = a[a00ci];
						a01d = a[a00di];
						a11a = a[a00ai];
						a11b = a[a00bi];
						a11c = a[a00ci];
						a11d = a[a00di];
						
						//XXX = XXX.add(t1.multy(a10));
						tmp[0] = w1[0]*a[a10ai] - w1[1]*a[a10bi];
						a00a = a00a + tmp[0];
						a10a = a10a - tmp[0];
						a01a = a01a + tmp[0];
						a11a = a11a - tmp[0];
						tmp[1] = w1[0]*a[a10bi] + w1[1]*a[a10ai];
						a00b = a00b + tmp[1];
						a10b = a10b - tmp[1];
						a01b = a01b + tmp[1];
						a11b = a11b - tmp[1];
						tmp[2] = w1[0]*a[a10ci] - w1[1]*a[a10di];
						a00c = a00c + tmp[2];
						a10c = a10c - tmp[2];
						a01c = a01c + tmp[2];
						a11c = a11c - tmp[2];
						tmp[3] = w1[0]*a[a10di] + w1[1]*a[a10ci];							
						a00d = a00d + tmp[3];
						a10d = a10d - tmp[3];
						a01d = a01d + tmp[3];
						a11d = a11d - tmp[3];
						
						//XXX = XXX.add(t2.multy(a01));
						tmp[0] = w2[0]*a[a01ai] - w2[2]*a[a01ci];
						a00a = a00a + tmp[0];
						a10a = a10a + tmp[0];
						a01a = a01a - tmp[0];
						a11a = a11a - tmp[0];
						tmp[1] = w2[0]*a[a01bi] - w2[2]*a[a01di];
						a00b = a00b + tmp[1];
						a10b = a10b + tmp[1];
						a01b = a01b - tmp[1];
						a11b = a11b - tmp[1];
						tmp[2] = w2[0]*a[a01ci] + w2[2]*a[a01ai];
						a00c = a00c + tmp[2];
						a10c = a10c + tmp[2];
						a01c = a01c - tmp[2];
						a11c = a11c - tmp[2];
						tmp[3] = w2[0]*a[a01di] + w2[2]*a[a01bi];	
						a00d = a00d + tmp[3];
						a10d = a10d + tmp[3];
						a01d = a01d - tmp[3];
						a11d = a11d - tmp[3];
						
						//XXX = XXX.add(t1.multy(t2.multy(a11)));
						tmp[0] = w2[0]*a[a11ai] - w2[2]*a[a11ci];
						tmp[1] = w2[0]*a[a11bi] - w2[2]*a[a11di];
						tmp[2] = w2[0]*a[a11ci] + w2[2]*a[a11ai];
						tmp[3] = w2[0]*a[a11di] + w2[2]*a[a11bi];	
						tmp2[0] = w1[0]*tmp[0] - w1[1]*tmp[1];
						tmp2[1] = w1[0]*tmp[1] + w1[1]*tmp[0];
						tmp2[2] = w1[0]*tmp[2] - w1[1]*tmp[3];
						tmp2[3] = w1[0]*tmp[3] + w1[1]*tmp[2];	
						a00a = a00a + tmp2[0];
						a00b = a00b + tmp2[1];
						a00c = a00c + tmp2[2];
						a00d = a00d + tmp2[3];
						a10a = a10a - tmp2[0];
						a10b = a10b - tmp2[1];
						a10c = a10c - tmp2[2];
						a10d = a10d - tmp2[3];
						a01a = a01a - tmp2[0];
						a01b = a01b - tmp2[1];
						a01c = a01c - tmp2[2];
						a01d = a01d - tmp2[3];
						a11a = a11a + tmp2[0];
						a11b = a11b + tmp2[1];
						a11c = a11c + tmp2[2];
						a11d = a11d + tmp2[3];
						
						a[a00ai] = a00a;
						a[a00bi] = a00b;
						a[a00ci] = a00c;
						a[a00di] = a00d;
						a[a10ai] = a10a;
						a[a10bi] = a10b;
						a[a10ci] = a10c;
						a[a10di] = a10d;
						a[a01ai] = a01a;
						a[a01bi] = a01b;
						a[a01ci] = a01c;
						a[a01di] = a01d;
						a[a11ai] = a11a;
						a[a11bi] = a11b;
						a[a11ci] = a11c;
						a[a11di] = a11d;
						
						//MIRROR PART START
						
						//mirr a00
						mirror(down, right, a, step2, N);
						
						//mirr a01
						mirror(down, right+step, a, step2, N);
												
						//mirr a10
						mirror(dsInd, right, a, step2, N);
						
						//mirr a11
						mirror(dsInd, right+step, a, step2, N);
								
						//MIRROR PART END
						
			
					}
				}
			}
		}
		return a;
	}

	
	private static float[] toVector(float[][] src) {
		float[] res = new float[src.length * src[0].length];
		int k = 0;
		for (int i = 0; i < src.length; i++) {
			for (int j = 0; j < src[0].length; j++) {
				res[k] = src[i][j];
				k++;
			}
		}
		return res;
	}
	
	private static int mirr(int x, int matrixWidth){
		int tmp = (x / matrixWidth) * matrixWidth * 2 - x;
		return (tmp == x) ? x : tmp + matrixWidth ;
	}
	
	private static void mirror(int downIndex, int rightIndex, float[] a, int step2, int N){
		int rMirr = mirr(rightIndex, step2);
		int dMirr = mirr(downIndex, step2);
		
		int a00ai = downIndex*N*4 + 4*rightIndex;
		int a00bi = a00ai + 1;
		int a00ci = a00ai + 2;
		int a00di = a00ai + 3;
		
		int aXXai = dMirr*N*4 + 4*rMirr;
		int aXXbi = aXXai + 1;
		int aXXci = aXXai + 2;
		int aXXdi = aXXai + 3;
		
		int a0Xai = downIndex*N*4 + 4*rMirr;
		int a0Xbi = a0Xai + 1;
		int a0Xci = a0Xai + 2;
		int a0Xdi = a0Xai + 3;
		
		int aX0ai = dMirr*N*4 + 4*rightIndex;
		int aX0bi = aX0ai + 1;
		int aX0ci = aX0ai + 2;
		int aX0di = aX0ai + 3;
		
		if (downIndex != dMirr && rightIndex != rMirr){
			a[aXXai] = a[a00ai];
			a[aXXbi] = -a[a00bi];
			a[aXXci] = -a[a00ci];
			a[aXXdi] = a[a00di];
		}
		if (downIndex != dMirr){
			a[aX0ai] = a[a00ai];
			a[aX0bi] = -a[a00bi];
			a[aX0ci] = a[a00ci];
			a[aX0di] = -a[a00di];
		}
		if (rightIndex != rMirr){
			a[a0Xai] = a[a00ai];
			a[a0Xbi] = a[a00bi];
			a[a0Xci] = -a[a00ci];
			a[a0Xdi] = -a[a00di];
		}
		
	}

	private static float[][] toMatrix(float[] src, int rows) {
		float[][] res = new float[rows][src.length / rows];
		int k = 0;
		for (int i = 0; i < rows; i++) {
			for (int j = 0; j < src.length / rows; j++) {
				res[i][j] = src[k];
				k++;
			}
		}
		return res;
	}
	
	private static void _inverse(float[][] src){
		for (int i = 0; i < src.length; i++){
			_inverse(src[i]);
		}
		float[] arr = new float[src.length];
		for (int i = 0; i < src[0].length; i++){
			for (int j = 0; j < src.length; j++){
				arr[j] = src[j][i];
			}
			inverse(arr);
			for (int j = 0; j < src.length; j++){
				src[j][i] = arr[j];
			}
		}
	}
	
	private static void _inverse(float[] arr){
		int arrlength = arr.length/2;
		int dl = arrlength/2;
		int st = arrlength - 1;
		int j = 0;
		for (int i = 0; i < st; i++){
			if (i < j){
				float s0 = arr[2*i];
				arr[2*i] = arr[2*j];
				arr[2*j] = s0;
				s0 = arr[2*i + 1];
				arr[2*i + 1] = arr[2*j + 1];
				arr[2*j + 1] = s0;
			}
			int k = dl;
			while (k <= j){
				j = j - k;
				k = k >> 1;
			}
			j = j + k;
		}
	}
	
	private static void __inverse(float[][] src){
		for (int i = 0; i < src.length; i++){
			__inverse(src[i]);
		}
		float[] arr = new float[src.length];
		for (int i = 0; i < src[0].length; i++){
			for (int j = 0; j < src.length; j++){
				arr[j] = src[j][i];
			}
			inverse(arr);
			for (int j = 0; j < src.length; j++){
				src[j][i] = arr[j];
			}
		}
	}
	
	private static void __inverse(float[] arr){
		int arrlength = arr.length/4;
		int dl = arrlength/2;
		int st = arrlength - 1;
		int j = 0;
		for (int i = 0; i < st; i++){
			if (i < j){
				float s0 = arr[4*i];
				arr[4*i] = arr[4*j];
				arr[4*j] = s0;
				
				s0 = arr[4*i + 1];
				arr[4*i + 1] = arr[4*j + 1];
				arr[4*j + 1] = s0;
				
				s0 = arr[4*i + 2];
				arr[4*i + 2] = arr[4*j + 2];
				arr[4*j + 2] = s0;
				
				s0 = arr[4*i + 3];
				arr[4*i + 3] = arr[4*j + 3];
				arr[4*j + 3] = s0;
			}
			int k = dl;
			while (k <= j){
				j = j - k;
				k = k >> 1;
			}
			j = j + k;
		}
	}
	
	private static void inverse(float[] arr){
		int dl = arr.length/2;
		int st = arr.length - 1;
		int j = 0;
		for (int i = 0; i < st; i++){
			if (i < j){
				float s0 = arr[i];
				arr[i] = arr[j];
				arr[j] = s0;
			}
			int k = dl;
			while (k <= j){
				j = j - k;
				k = k >> 1;
			}
			j = j + k;
		}
	}
	
	private static float[][] transformToComplexMatrix(float[][] src){
		float[][] res = new float[src.length][src[0].length*2];
		for (int i = 0; i < src.length; i++){
			for (int j = 0; j < src.length; j++){
				res[i][j*2] = src[i][j];
			}
		}
		return res;
	}
	
	private static float[][] transformToGyperComplexMatrix(float[][] src){
		float[][] res = new float[src.length][src[0].length*4];
		for (int i = 0; i < src.length; i++){
			for (int j = 0; j < src.length; j++){
				res[i][j*4] = src[i][j];
			}
		}
		return res;
	}

	public static void main(String args[]) throws IOException, CudaCRuntimeException {
		float[][] arr = Generator.generateSimpleMatrix(8);
		arr = transformToGyperComplexMatrix(arr);
		
		//String moduleName = CudaCompiler.preparePtxFile(FFT_MIRROR_MODULE_SRC, true, " -DWLIST_SIZE=2048");
		//CUfunction function = CudaCompiler.getFunction(FFT_FUNCTION_NAME, moduleName);
		long time = System.nanoTime();
		//arr = fftGCMirror(arr, function);
		arr = __fft(arr);
		time = System.nanoTime() - time;
		Printer.printVector(arr);
		System.out.println(time);
		
	}
}